use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use anyhow::Result;

use protoactor_rs::actor::{Actor, ActorContext, ActorError, DefaultActorContext, PID, Props};
use protoactor_rs::system::ActorSystem;
use protoactor_rs::remote::{Remote, RemoteConfig};
use protoactor_rs::cluster::{
    Grain, GrainContext, GrainState, Member, MemberStatus,
    ClusterSystem, ClusterConfig, ClusterEvent,
};

// 定义一个简单的虚拟Actor
struct HelloGrain {
    id: String,
    counter: usize,
}

// 虚拟Actor的消息
struct Hello {
    name: String,
}

// 虚拟Actor的响应
#[derive(Clone, Debug)]
struct HelloResponse {
    message: String,
    count: usize,
}

impl HelloGrain {
    fn new(id: &str) -> Self {
        Self {
            id: id.to_string(),
            counter: 0,
        }
    }
}

#[async_trait]
impl Actor for HelloGrain {
    type Context = GrainContext;
    
    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn std::any::Any + Send>) -> Result<(), ActorError> {
        // 更新活动状态
        ctx.update_activity();
        
        if let Some(hello) = msg.downcast_ref::<Hello>() {
            let mut grain = self.clone();
            grain.counter += 1;
            
            let response = HelloResponse {
                message: format!("Hello, {}! (from grain {})", hello.name, self.id),
                count: grain.counter,
            };
            
            ctx.respond(response);
        } else if let Some(cmd) = msg.downcast_ref::<&str>() {
            if *cmd == "get_state" {
                ctx.respond(GrainState::Active);
            }
        }
        
        Ok(())
    }
}

#[async_trait]
impl Grain for HelloGrain {
    fn id(&self) -> &str {
        &self.id
    }
    
    fn kind(&self) -> &str {
        "hello_grain"
    }
    
    async fn activate(&mut self) -> Result<()> {
        println!("Activating grain: {}", self.id);
        Ok(())
    }
    
    async fn deactivate(&mut self) -> Result<()> {
        println!("Deactivating grain: {}", self.id);
        Ok(())
    }
}

impl Clone for HelloGrain {
    fn clone(&self) -> Self {
        Self {
            id: self.id.clone(),
            counter: self.counter,
        }
    }
}

// 假设ClusterSystem还没有实现spawn_grain方法，我们这里扩展一下
trait GrainExtensions {
    async fn spawn_grain(&self, kind: &str, id: &str) -> Result<PID>;
    async fn send_to_grain(&self, kind: &str, id: &str, msg: impl std::any::Any + Send + 'static) -> Result<()>;
    async fn request_from_grain<R: std::any::Any + Send + 'static>(
        &self, kind: &str, id: &str, msg: impl std::any::Any + Send + 'static
    ) -> Result<R>;
}

impl GrainExtensions for ClusterSystem {
    async fn spawn_grain(&self, kind: &str, id: &str) -> Result<PID> {
        // 先获取角色为kind的节点
        let nodes = self.get_members_by_role(kind).await;
        
        if nodes.is_empty() {
            // 如果没有适合的节点，就使用本地节点
            let props = Props::new_fn(move || Box::new(HelloGrain::new(id)));
            
            // 使用本地系统创建
            let system = ActorSystem::new();
            let pid = system.spawn_named(props, &format!("{}_{}",kind, id)).await?;
            Ok(pid)
        } else {
            // 选择一个节点（简单地使用第一个）
            let node = &nodes[0];
            let address = format!("{}:{}", node.host, node.port);
            
            // 远程创建Actor
            let pid = self.spawn_remote::<HelloGrain>(
                &address, 
                kind, 
                id, 
                Duration::from_secs(5),
            ).await?;
            
            Ok(pid)
        }
    }
    
    async fn send_to_grain(&self, kind: &str, id: &str, msg: impl std::any::Any + Send + 'static) -> Result<()> {
        // 生成grain的ID
        let grain_id = format!("{}_{}",kind, id);
        
        // 查找或创建grain
        let pid = self.spawn_grain(kind, id).await?;
        
        // 发送消息
        self.send(&pid, msg).await?;
        
        Ok(())
    }
    
    async fn request_from_grain<R: std::any::Any + Send + 'static>(
        &self, kind: &str, id: &str, msg: impl std::any::Any + Send + 'static
    ) -> Result<R> {
        // 生成grain的ID
        let grain_id = format!("{}_{}",kind, id);
        
        // 查找或创建grain
        let pid = self.spawn_grain(kind, id).await?;
        
        // 发送请求并等待响应
        let response = self.request::<R>(&pid, msg, Duration::from_secs(5)).await?;
        
        Ok(response)
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // 初始化日志
    tracing_subscriber::fmt::init();
    
    // 创建actor系统
    let system = ActorSystem::new();
    
    // 创建远程配置
    let remote_config = protoactor_rs::remote::Config::new()
        .with_host("127.0.0.1:8090")
        .build();
    
    // 创建远程处理模块
    let mut remote = Remote::new(system.clone(), remote_config);
    remote.start().await?;
    
    // 注册HelloGrain
    remote.register::<HelloGrain>("hello_grain", Props::new_fn(|| Box::new(HelloGrain::new("temp"))))?;
    
    // 创建集群配置
    let cluster_config = ClusterConfig {
        name: "hello_grain".to_string(),  // 设置角色为hello_grain
        host: "127.0.0.1".to_string(),
        port: 8090,
        heartbeat_interval: 1,
        unreachable_timeout: 5,
        down_timeout: 10,
    };
    
    // 创建集群系统
    let (mut cluster, mut cluster_events) = ClusterSystem::new_with_memory_provider(
        cluster_config, 
        Arc::new(remote)
    );
    
    // 启动集群系统
    cluster.start().await?;
    
    // 监听集群事件（在实际使用中，应该在单独的任务中处理）
    tokio::spawn(async move {
        while let Some(event) = cluster_events.recv().await {
            println!("Cluster event: {:?}", event);
        }
    });
    
    // 创建克隆用于多个任务
    let cluster_clone = Arc::new(cluster);
    
    // 创建一个任务来使用虚拟Actor
    let task_cluster = cluster_clone.clone();
    tokio::spawn(async move {
        // 等待集群初始化
        tokio::time::sleep(Duration::from_secs(1)).await;
        
        for i in 1..=5 {
            let name = format!("User {}", i);
            
            // 发送请求到grain并接收响应
            match task_cluster.request_from_grain::<HelloResponse>(
                "hello_grain", 
                "grain1", 
                Hello { name },
            ).await {
                Ok(response) => {
                    println!("Got response: {:?}", response);
                },
                Err(e) => {
                    println!("Error: {}", e);
                }
            }
            
            // 稍微等待
            tokio::time::sleep(Duration::from_millis(500)).await;
        }
    });
    
    // 保持主线程运行
    tokio::time::sleep(Duration::from_secs(10)).await;
    
    // 停止集群
    let mut cluster = Arc::try_unwrap(cluster_clone).unwrap();
    cluster.stop().await?;
    
    println!("Done!");
    Ok(())
} 